
import numpy as np

try:
    import cupyx
except ImportError:
    cupyx = np

from core.cuda import cuda_module
from core.function import Function
import utils

# 切片
class GetItemGrad(Function):
    def __init__(self, slices, in_shape):
        """
        初始化 GetItemGrad 函数。

        Parameters:
            slices: 切片信息。
            in_shape: 输入数据的形状。

        Returns:
            None
        """
        self.slices = slices
        self.in_shape = in_shape

    def forward(self, gy):
        """
        前向传播函数，实现对相对于切片结果的梯度计算。

        Parameters:
            gy: 相对于前向传播结果的梯度。

        Returns:
            gx: 相对于输入数据的梯度。
        """
        xp = cuda_module
        gx = xp.zeros(self.in_shape, dtype=gy.dtype)

        # np.add.at(gx, self.slices, gy)

        if xp is np:
            np.add.at(gx, self.slices, gy)
        else:
            cupyx.scatter_add(gx, self.slices, gy)
        return gx

    def backward(self, ggx):
        """
        反向传播函数，计算相对于输入数据的梯度。

        Parameters:
            ggx: 相对于前向传播结果的梯度。

        Returns:
            ggy: 相对于输入数据的梯度。
        """

        return utils.functions_collect.get_item(ggx, self.slices)


if __name__ == '__main__':
    pass
